package lesson4.example8;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.expr.ExprEditor;
import javassist.expr.MethodCall;

import java.lang.instrument.ClassFileTransformer;
import java.security.ProtectionDomain;

public class MyTransformer implements ClassFileTransformer {

    @Override
    public byte[] transform(
		ClassLoader       classLoader, 
		String            className, 
		Class<?>          classBeingRedefined, 
		ProtectionDomain  protectionDomain, 
		byte[]            bytes
    ) {
        try {
            if (!className.contains("lesson4")) {
                return null;
            }
            String currentClassName = className.replaceAll("/", ".");
            CtClass ctClass = ClassPool.getDefault().get(currentClassName);
            CtBehavior[] methods = ctClass.getDeclaredBehaviors();
            for (CtBehavior method : methods) {
                enhanceMethod(method);
            }
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        
        return null;
    }

    private void enhanceMethod(CtBehavior method) throws Exception {
        String methodName = method.getName();
        String s  = 
        		"{ \n" + 
        		"  long start = System.nanoTime();\n" +
                "  $_ = $proceed($$);\n"+
                "  System.out.println(\"method:"+methodName+" cost:\" +(System.nanoTime() - start)+ \"ns\"); \n"+
                "} \n";

        // Modifies the method/constructor body.
        method.instrument(new ExprEditor() {
            @Override
            public void edit(MethodCall methodCall) throws CannotCompileException {
                methodCall.replace(s);
            }
        });
    }

}
